#!/usr/bin/env python3

'''
    Software License

    Copyright (C) 2021-05-24  Xoronos

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, version 3.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
'''

'''
    Liabilities

    The software is provided "AS IS" without any warranty of any kind, either expressed,
    implied, or statutory, including, but not limited to, any warranty that the software
    will conform to specifications, any implied warranties of merchantability, fitness
    for a particular purpose, and freedom from infringement, and any warranty that the
    documentation will conform to the software, or any warranty that the software will
    be error free.

    In no event shall Xoronos be liable for any damages, including, but not limited to,
    direct, indirect, special or consequential damages, arising out of, resulting from,
    or in any way connected with this software, whether or not based upon warranty,
    contract, tort, or otherwise, whether or not injury was sustained by persons or
    property or otherwise, and whether or not loss was sustained from, or arose out of
    the results of, or use of, the software or services provided hereunder.
    
    To request the provided software under a different license you can contact us at
    support@xoronos.com
'''

from PIL import Image
import numpy as np
import pickle
import soundfile as sf
import mutagen.flac 
import zlib
import struct
import sys
import json
import os

if not os.path.isdir('/usr/lib/xoronos/py/') :
    print("directory /usr/lib/xoronos/py/ does not exist")
    exit(0)

sys.path.append('/usr/lib/xoronos/py/')

from splitjoin import msb_lsb_split
from splitjoin import msb_lsb_join

PNG = 0
FLAC = 1

def save_array_metadata_png(metadata_file, array, ftype):

    # Serialize the metadata (shape, data type, etc.)
    metadata = {
        'shape': array.shape,
        'size': array.size,
        'dtype': array.dtype
    }

    # Serialize metadata to bytes
    metadata_bytes = pickle.dumps(metadata)

    crc32 = zlib.crc32( struct.pack('<c', bytes([ftype])) + metadata_bytes )

    with open(metadata_file, 'wb') as f:
        f.write(struct.pack('<c', bytes([ftype])))
        f.write(metadata_bytes)
        f.write(struct.pack('<I',crc32))

def save_array_metadata_flac(metadata_file, flac_info, sample_rate, array, ftype, image_array ):
    # Serialize the metadata (shape, data type, etc.)
    metadata = {
        'shape': array.shape,
        'size': array.size,
        'dtype': array.dtype,
        'sample_rate': sample_rate,
        'image_array': image_array,
        'flac_info': flac_info
    }

    # Serialize metadata to bytes
    metadata_bytes = pickle.dumps(metadata)

    crc32 = zlib.crc32(struct.pack('<c', bytes([ftype])) + metadata_bytes)

    with open(metadata_file, 'wb') as f:
        f.write(struct.pack('<c', bytes([ftype])))
        f.write(metadata_bytes)
        f.write(struct.pack('<I', crc32))

def load_array_metadata(metadata_file, extended_file_extension):

    # Deserialize the metadata
    with open(metadata_file, 'rb') as f:
        binary_data = f.read()

    # Extract stored ftype value
    stored_ftype = struct.unpack('<c', binary_data[0:1])[0]
    if ( int.from_bytes(stored_ftype,"little") != extended_file_extension ) :
        print('error in expected file type stored_ftype {} extended_file_extension {}'.format(stored_ftype , extended_file_extension))
        exit(-1)

    # Extract stored crc32 value
    stored_crc32 = struct.unpack('<I', binary_data[-4:])[0]

    # Extract stored metadata
    metadata_bytes = binary_data[1:-4]

    # Verify crc32
    calculated_crc32 = zlib.crc32(struct.pack('<c',stored_ftype) + metadata_bytes)

    if calculated_crc32 != stored_crc32 :
        print('CRC32 error: stored_crc32 {}, calculated_crc32 {}'.format(stored_crc32, calculated_crc32))
        exit(-1)
    else :
        return pickle.loads(metadata_bytes) , ord(stored_ftype)


def save_array_data(msbs_file, lsbs_file, bitsize, bit_skip, array):

    msbs_array , lsbs_array = msb_lsb_split( array , bitsize, bit_skip )

    # Create a sample NumPy array
    msbs_byte_array = bytearray(np.array(msbs_array, dtype=np.uint8).tobytes())
    lsbs_byte_array = bytearray(np.array(lsbs_array, dtype=np.uint8).tobytes())

    # Serialize the msbs array data
    with open(msbs_file, 'wb') as f:
        f.write( msbs_byte_array )

    # Serialize the lsbs array data
    with open(lsbs_file, 'wb') as f:
        f.write( lsbs_byte_array )


def load_array_data( msbs_file, lsbs_file, bitsize, bit_skip, num_el ):

    # Deserialize the msbs data
    with open(msbs_file, 'rb') as f:
        msbs_array_byte_array = bytearray(f.read( ))
        msbs_array = np.frombuffer( msbs_array_byte_array, dtype=np.uint8 )
 
    # Deserialize the msbs data
    with open(lsbs_file, 'rb') as f:
        lsbs_array_byte_array = bytearray(f.read( ))
        lsbs_array = np.frombuffer( lsbs_array_byte_array, dtype=np.uint8 )

    array_data = msb_lsb_join( msbs_array.astype(np.uint8) , lsbs_array.astype(np.uint8) , num_el , bitsize , bit_skip )

    return array_data


def reconstruct_array(metadata, data):
    # Reconstruct the numpy array using the deserialized metadata and data
    array = np.ndarray(metadata['shape'], dtype=metadata['dtype'], buffer=data)
    return array

# Create an array from the four matrices

def create_array_from_matrices(red_matrix, green_matrix, blue_matrix, alpha_matrix=None):
    if alpha_matrix is not None:
        img_array = np.dstack((red_matrix, green_matrix, blue_matrix, alpha_matrix))
    else:
        img_array = np.dstack((red_matrix, green_matrix, blue_matrix))
    return img_array

# Split an array into the four matrices

def split_array_into_matrices(img_array):
    red_matrix = img_array[:, :, 0]
    green_matrix = img_array[:, :, 1]
    blue_matrix = img_array[:, :, 2]
    if img_array.shape[2] == 4:  # If array has alpha channel
        alpha_matrix = img_array[:, :, 3]
        return red_matrix, green_matrix, blue_matrix, alpha_matrix
    else:
        return red_matrix, green_matrix, blue_matrix, None

def png_to_three_arrays(input_file, metadata_file, msbs_file, lsbs_file):
    # Open the PNG image
    img = Image.open(input_file)
    
    # Convert the image to a numpy array
    img_array = np.array(img)
    
    if img.mode == 'RGBA' :
        # Separate the RGBA
        red_channel = img_array[:,:,0]
        green_channel = img_array[:,:,1]
        blue_channel = img_array[:,:,2]
        alpha_channel = img_array[:,:,3]
        out_array = create_array_from_matrices ( red_channel, green_channel, blue_channel, alpha_channel )

    elif img.mode == 'RGB':
        # Separate the RGB
        red_channel = img_array[:,:,0]
        green_channel = img_array[:,:,1]
        blue_channel = img_array[:,:,2]
        out_array = create_array_from_matrices ( red_channel, green_channel, blue_channel, None )

    else:
        # If the image does not have colors
        out_array = np.array(img_array)
    
    ftype = PNG
    save_array_metadata_png(metadata_file, out_array, ftype )
    bitsize = out_array.dtype.itemsize * 8
    bit_skip = 0
    save_array_data(msbs_file, lsbs_file, bitsize, bit_skip, out_array.reshape(-1).astype(np.uint64))



def two_arrays_to_png( metadata, data, output_file):
    bitsize = metadata['dtype'].itemsize * 8
    num_el = metadata['size']
    shape = metadata['shape']
    img_array = reconstruct_array(metadata, data)
    
    if len(shape) == 3 :

        red_matrix, green_matrix, blue_matrix, alpha_matrix = split_array_into_matrices(img_array)
        
        # Merge the matrices into a single array if an alpha matrix is provided
        if alpha_matrix is not None:
            img_array = np.dstack((red_matrix, green_matrix, blue_matrix, alpha_matrix))
        else:
            img_array = np.dstack((red_matrix, green_matrix, blue_matrix))

    else :
        img_array = np.array(img_array)

    
    # Create an image object from the array
    if bitsize == 8 :
        img = Image.fromarray(img_array.astype('uint8'))
    elif bitsize == 16 :
        img = Image.fromarray(img_array.astype('uint16'))
    
    # Save the image to a PNG file
    img.save(output_file)


def flac_to_three_arrays(input_file, metadata_file, msbs_file, lsbs_file):
    out_array, sample_rate = sf.read(input_file)
    flac_info = mutagen.flac.FLAC(input_file)

    image_array = []
    for picture in flac_info.pictures :
        image_array.append(picture)

    ftype = FLAC

    bitsize = out_array.dtype.itemsize * 8
    save_array_metadata_flac(metadata_file, flac_info, sample_rate, out_array, ftype, image_array )

    bit_skip = 0
    if out_array.dtype == np.float64 :
        out_array = np.float64(out_array).view(np.uint64)
        bit_skip = 12
    elif out_array.dtype == np.float32 :
        out_array = np.float32(out_array).view(np.uint32)
        bit_skip = 9
    out_array = out_array.reshape(-1).astype(np.uint64)

    save_array_data(msbs_file, lsbs_file, bitsize, bit_skip, out_array)


def two_arrays_to_flac(metadata, data, output_file):
    # Write audio data to FLAC file
    sf.write(output_file, data, metadata['sample_rate'])

    # Update FLAC metadata blocks
    flac_file = mutagen.flac.FLAC(output_file)
    flac_file.clear()

    # Copy metadata from the source FLAC file to the target FLAC file
    for key, value in metadata['flac_info'].items():
        flac_file[key] = value

    for picture in metadata['image_array'] :
        flac_file.add_picture(picture)

    flac_file.save()


def three_arrays_to_file( metadata_file, msbs_file, lsbs_file, output_file, file_extension):
    metadata, ftype = load_array_metadata ( metadata_file, file_extension)
    bitsize = metadata['dtype'].itemsize * 8
    num_el = metadata['size']

    if ( metadata['dtype'].type == np.float64 ) :
        bit_skip = 12
        data = load_array_data ( msbs_file, lsbs_file, bitsize, bit_skip, num_el )
        data = np.uint64(data).view(np.float64)
        data = data.reshape(metadata['shape']).astype(metadata['dtype'].type)
    elif ( metadata['dtype'].type == np.float32 ) :
        bit_skip = 9
        data = load_array_data ( msbs_file, lsbs_file, bitsize, bit_skip, num_el )
        data = np.uint64(data).view(np.float64)
    else :
        bit_skip = 0
        data = load_array_data ( msbs_file, lsbs_file, bitsize, bit_skip, num_el )
        data = data.reshape(metadata['shape']).astype(metadata['dtype'].type)

    if ftype == PNG :
        two_arrays_to_png( metadata, data, output_file )
    elif ftype == FLAC :
        two_arrays_to_flac( metadata, data, output_file )

################
# Parser lexer #
################

tool_name_str = "xrnconv-cli"
split_str = "--std2xrn"
join_str = "--xrn2std"
source_str = "-source"
destination_str = "-destination"
msbs_str = "-msbs"
lsbs_str = "-lsbs"
metadata_str = "-metadata"

list_of_commands = [ split_str, join_str ]
list_of_options = [ source_str, destination_str , msbs_str, lsbs_str, metadata_str ]

list_of_everything = list_of_commands + list_of_options
arglist = [0 for i in range(len(list_of_commands) + len(list_of_options))]

def print_help ():
   print("help page for the xrnconv-cli command\n")
   print("to split a standard file in three files")
   print("{} {} {} srcfile {} msbfile {} lsbfile {} matadatafile ".format( tool_name_str, split_str, source_str, msbs_str, lsbs_str, metadata_str))
   print("\nto join three files in a standard file")
   print("{} {} {} dstfile {} msbfile {} lsbfile {} matadatafile ".format( tool_name_str, join_str, destination_str, msbs_str, lsbs_str, metadata_str))

if ( len(sys.argv) == 2 ) and ( sys.argv[1] == "--help" ) :
   print_help ()
   exit (-1)

if len(sys.argv) < 2 :
   print ("no command specified")
   print_help ()
   exit (-1)


followed_by_a_file = 0
argrecognized = 0
prevarg = ""

input_file = ""
metadata_file = ""
msbs_file = ""
lsbs_file = ""
output_file = ""
jsonstr = ""

# parse all the commands
for arg in sys.argv[1:]:

   if ( arg not in list_of_commands ) and ( arg not in list_of_options ) :
       argrecognized = 0
   else :
       argrecognized = 1

   # check command/option validity
   if followed_by_a_file == 0 :
       if argrecognized == 0:
           print ("argument {} not recognized".format(arg))
           print_help ()
           exit (-1)

       arglist[ list_of_everything.index(arg) ] += 1
       if arglist[ list_of_everything.index(arg) ] > 1 :
           print ("command {} repeted too many times".format(arg))
           print_help ()
           exit (-1)

       if ( arg == source_str ) or ( arg == destination_str ) or ( arg == msbs_str ) or ( arg == lsbs_str ) or ( arg == metadata_str ) :
           followed_by_a_file = 1

   else :

       if followed_by_a_file == 1 :
           if argrecognized == 1 :
               print ("argument {} should be followed by a file".format(prevarg))
               print_help ()
               exit (-1)
           else :
               if prevarg == msbs_str :
                   msbs_file = arg
               elif prevarg == lsbs_str :
                   lsbs_file = arg
               elif prevarg == metadata_str :
                   metadata_file = arg
               elif prevarg == source_str :
                   input_file = arg
               elif prevarg == destination_str :
                   output_file = arg


       followed_by_a_file = 0

   prevarg = arg

# common checks
if followed_by_a_file == 1 :
    print ("string or file not specified")
    print_help ()
    exit (-1)

if ( ( arglist[ list_of_everything.index(split_str)] == 0 ) and \
     ( arglist[ list_of_everything.index(join_str)] == 0 )) :
    print ("at least on of the following commands should be specified {} {}".format( join_str, split_str))
    print_help ()
    exit (-1)
if ( arglist[ list_of_everything.index(split_str)] == 1 ) and \
   ( arglist[ list_of_everything.index(join_str)] == 1 ) :
    print ("argument {} should not be specified whith {}".format(split_str,join_str))
    print_help ()
    exit (-1)
if ( arglist[ list_of_everything.index(destination_str)] == 1 ) and \
   ( arglist[ list_of_everything.index(source_str)] == 1 ) :
    print ("argument {} should not be specified whith {}".format(destination_str,source_str))
    print_help ()
    exit (-1)
if ( arglist[ list_of_everything.index(split_str)] == 1 ) and \
   (( arglist[ list_of_everything.index(msbs_str)] == 0 ) or \
   ( arglist[ list_of_everything.index(lsbs_str)] == 0 ) or \
   ( arglist[ list_of_everything.index(metadata_str)] == 0 ) or \
   ( arglist[ list_of_everything.index(source_str)] == 0 )) :
    print ("wrong {} command".format(split_str))
    print_help ()
    exit (-1)
if ( arglist[ list_of_everything.index(join_str)] == 1 ) and \
   (( arglist[ list_of_everything.index(msbs_str)] == 0 ) or \
   ( arglist[ list_of_everything.index(lsbs_str)] == 0 ) or \
   ( arglist[ list_of_everything.index(metadata_str)] == 0 ) or \
   ( arglist[ list_of_everything.index(destination_str)] == 0 )) :
    print ("wrong {} command".format(join_str))
    print_help ()
    exit (-1)

###############
# Run command #
###############

if ( arglist[ list_of_everything.index(split_str)] == 1 ) :
   _ , file_extension = os.path.splitext(input_file)
   if ( file_extension == ".png" ) or ( file_extension == ".PNG" ) :
       png_to_three_arrays(input_file, metadata_file, msbs_file, lsbs_file)
   if ( file_extension == ".flac" ) or ( file_extension == ".FLAC" ) :
       flac_to_three_arrays(input_file, metadata_file, msbs_file, lsbs_file)

if ( arglist[ list_of_everything.index(join_str)] == 1 ) :
   _ , file_extension_str = os.path.splitext(output_file)
   if ( file_extension_str == ".png" ) or ( file_extension_str == ".PNG" ) :
       file_extension = PNG
   elif ( file_extension_str == ".flac" ) or ( file_extension_str == ".FLAC" ) :
       file_extension = FLAC
   three_arrays_to_file( metadata_file, msbs_file, lsbs_file, output_file, file_extension )


